from copy import deepcopy
import numpy as np
import pandas as pd
import sys


'''
In this problem you write your own K-Means
Clustering code.

Your code should return a 2d array containing
the centers.

'''
# Import the dataset
data = []
with open('data/iris.data', 'r') as file:
    for line in file:
        data.append(line.strip().split(','))
data = [list(map(float, line[:4])) for line in data[1:]]
data = np.array(data)

# Make 3  clusters
k = 3
# Initial Centroids
C = [[2.,  0.,  3.,  4.], [1.,  2.,  1.,  3.], [0., 2.,  1.,  0.]]
C = np.array(C)
print("Initial Centers")
print(C)


def get_assignment(data, centers):
    """Get the label of each data point."""
    assign = np.zeros(data.shape[0])
    for i in range(data.shape[0]):
        assign[i] = np.argmin(get_distance(data[i], centers, 'l2_vector'))
    return assign


def get_data_in_cluster(labels, data, target):
    """Find datapoints assigned to the targeted cluster."""
    return data[labels==target]


def get_distance(a, b, method):
    """Compute two kinds of distance. 
        1. Distance between one instance (a) and every center in center array (b), 
           use 'l2_vector' method, return an array of (k, #dimensions).
        2. Absolute difference between old (a) and new (b) centers in two 
           iterations, use 'l2_matrix' method, return a scalar.
    """
    if method == 'l2_vector':
        return np.linalg.norm(a - b, axis=1) ** 2
    # if method == 'l1':
    #     return np.absolute(a - b).sum()
    if method == 'l2_matrix':
        return np.linalg.norm(a - b)


def k_means(C):
    """Perform k-means clustering using Lloyd's algorithm.

    Args:
        C: Cluster centers in shape (k, #dimensions).
        data: input data, (#instances, #dimensions).
        labels: assignment of each instance, (#instances,).
    Returns:
        C: Final cluster centers (k, #dimensions).
    """
    # Make sure C is an array.
    C = np.array(C)
    # Compute initial assignment.
    labels = get_assignment(data, C)
    # Initialize diff to be the max integer.
    diff = sys.maxsize
    
    # Loop until convergence.
    while diff > 10e-3:
        # Re-center.
        C_old = deepcopy(C)
        for target_k in range(k):
            cluster_data = get_data_in_cluster(labels, data, target_k)
            C[target_k] = np.mean(cluster_data, axis=0) 
        # Re-assign.
        labels = get_assignment(data, C)
        # Compute difference between old and new centers.
        diff = get_distance(C_old, C, 'l2_matrix')

    # return C_final
    return C

print("Final centers:")
print(k_means(C))






